# view_trajectory.py
# 用于查看保存在 trajectories.h5 中的游戏轨迹

import cv2
import h5py
import numpy as np
from config import TRAJECTORY_SAVE_PATH

def list_episodes(file_path):
    """ 列出 HDF5 文件中所有 episode """
    with h5py.File(file_path, 'r') as f:
        episodes = [key for key in f.keys() if key.startswith("episode_")]
        episodes.sort(key=lambda x: int(x.split("_")[1]))  # 按编号排序
        return episodes

def get_step_range():
    """ 获取用户输入的 step 范围 """
    try:
        start = int(input("请输入起始 step (如 10): "))
        end   = int(input("请输入结束 step (如 20): "))
        if start < 0 or end < start:
            raise ValueError
        return start, end
    except:
        print("输入无效，使用默认范围: 0 到 10")
        return 0, 10

def display_trajectory_step(step_group):
    """ 显示单个 step 的信息和原始图像（假设数据干净） """
    count       = step_group["count"][()]
    frame_stack = step_group["frame"][()]  # uint8 shape: (T, H, W)
    reward      = step_group["reward"][()]
    knight_hp   = step_group["knight_hp"][()]
    boss_hp     = step_group["boss_hp"][()]
    action      = step_group["action"][()]

    print(f"Step {int(count):4d} | "
          f"Knight HP: {knight_hp:2d} | "
          f"Boss HP: {boss_hp:.3f} | "
          f"Reward: {reward:+.1f} | "
          f"Action: {action}")
    # print(f"🖼️  Frame shape: {frame_stack.shape}, dtype: {frame_stack.dtype}")

    # 拼接所有帧（水平方向）
    concatenated = np.hstack(frame_stack)  # shape: (100, 256 * T)

    # 放大图像
    # enlarged_frame = cv2.resize(concatenated, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)

    cv2.imshow("Raw Frame (Unmodified)", concatenated)
    key = cv2.waitKey(0)
    if key in [ord('q'), 27]:
        cv2.destroyAllWindows()
        exit()

def view_trajectory_interactive(file_path="trajectories.h5"):
    """ 交互式查看轨迹 """
    episodes = list_episodes(file_path)
    if not episodes:
        print("❌ 文件中没有找到任何 episode 数据！")
        return

    print("📁 可用的 episode:")
    for i, ep in enumerate(episodes):
        with h5py.File(file_path, 'r') as f:
            num_steps = len(f[ep])
        print(f"  {i}: {ep} ({num_steps} steps)")

    try:
        choice = int(input(f"\n请选择 episode 编号 [0-{len(episodes)-1}]: "))
        selected_episode = episodes[choice]
    except (ValueError, IndexError):
        print("输入无效，选择第一个 episode")
        selected_episode = episodes[0]

    start, end = get_step_range()

    print(f"\n🔍 正在加载 {selected_episode} 的 step {start} 到 {end}...\n")

    # 打开 HDF5 文件，遍历指定范围的 step
    with h5py.File(file_path, 'r') as f:
        episode_grp = f[selected_episode]
        found_any = False
        for step_idx in range(start, end + 1):
            step_name = f"step_{step_idx}"
            if step_name in episode_grp:
                found_any = True
                step_group = episode_grp[step_name]
                display_trajectory_step(step_group)

        if not found_any:
            print(f"⚠️ 在 {selected_episode} 中未找到 step_{start} 到 step_{end} 的数据。")

    print("\n🔚 查看结束。关闭图像窗口...")
    cv2.waitKey(0)  # 等待用户按键
    cv2.destroyAllWindows()

if __name__ == "__main__":
    view_trajectory_interactive(TRAJECTORY_SAVE_PATH)